package service

import (
	"bufio"
	"bytes"
	"context"
	"github.com/gin-gonic/gin"
	"github.com/goccy/go-json"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"strings"
)

type ChatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// WebSocket 升级器
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// WebSocket 处理入口
func WsHandler(c *gin.Context) {
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Println("WebSocket 连接失败:", err)
		return
	}
	defer conn.Close()
	handleOllamaStream(conn)
}

func handleOllamaStream(conn *websocket.Conn) {
	var history []ChatMessage

	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			log.Println("读取消息失败:", err)
			break
		}

		// 每次收到用户输入
		history = append(history, ChatMessage{
			Role:    "user",
			Content: string(msg),
		})

		// 发起请求，获取 AI 回复（带上整个历史）
		aiReply, err := queryOllama(history, conn)
		if err != nil {
			log.Println("获取模型响应失败:", err)
			break
		}

		// 将模型回复也写入历史
		history = append(history, ChatMessage{
			Role:    "assistant",
			Content: aiReply,
		})
	}
}

func queryOllama(messages []ChatMessage, conn *websocket.Conn) (string, error) {
	ctx := context.Background()
	body := map[string]interface{}{
		"model":    "Qwen2.5-3B-Int4",
		"stream":   true,
		"messages": messages,
	}

	jsonData, _ := json.Marshal(body)
	req, _ := http.NewRequestWithContext(ctx, "POST", "http://10.11.203.245:6300/v1/chat/completions", bytes.NewBuffer(jsonData))
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	var fullResponse strings.Builder
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || line == "data: [DONE]" {
			continue
		}

		raw := strings.TrimPrefix(line, "data: ")

		var chunk struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
		}

		if err := json.Unmarshal([]byte(raw), &chunk); err == nil {
			if len(chunk.Choices) > 0 {
				content := chunk.Choices[0].Delta.Content
				if content != "" {
					fullResponse.WriteString(content)
					conn.WriteMessage(websocket.TextMessage, []byte(content))
				}
			}
		}
	}

	return fullResponse.String(), nil
}
